from sklearn import metrics, svm
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import precision_recall_fscore_support as score
from sklearn.neighbors import KNeighborsClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.naive_bayes import GaussianNB
from sklearn.ensemble import GradientBoostingClassifier

def create_model(name):
    """创建机器学习模型

    Args:
        name (str): ["lr", "knn", "svm",  "nb", "gbt"]

    Returns:
        model: 机器学习模型
    """
    if name == 'svm':
        model = svm.LinearSVC(max_iter=1200, dual=False)
    elif name == 'knn':
        model = KNeighborsClassifier(n_neighbors=10)
    elif name == 'lr':
        model = LogisticRegression(max_iter=1200)   
    elif name == 'nb':
        model = GaussianNB(var_smoothing=1.0)
    elif name == 'gbt':
        model = GradientBoostingClassifier(n_estimators=100, 
                                          learning_rate=1.0, 
                                          max_depth=1,
                                          random_state=0)
    else:
        print('请输入正确的机器学习模型')
    return model


def get_model_metrics(test_output, test_predict):
    """返回机器学习训练结果

    Args:
        test_output (1d array-like): 实际输出
        test_predict (1d array-like): 预测结果

    Returns:
        acc, precision, recall, f1_score
    """    
    precision, recall, f1_score, _ = score(
        test_output, test_predict, average='macro')
    acc = metrics.accuracy_score(test_output, test_predict)
    return acc, precision, recall, f1_score